######################################################################
#Custom functions for replication
######################################################################

############################
#Make equations
############################
make.equation <- function(dvs, ivs, interactions = c()){
  return(paste(dvs, "~", paste(c(ivs, interactions), sep = "", collapse = " + ")))
}

############################
#Make formula for feglm
############################
make.formula = function(dvs, ivs, interactions = c()){
  return(paste(dvs, "~", paste(c(ivs, interactions), sep = "", collapse = " + ")) %>%
            paste(., "| iso2c + year | iso2c", sep = " ") %>%
            as.formula())
}


######################################
#Estimate bpCausal with a vector of DVs
  #With the following parameters
######################################
# random effect: choose from ("unit", "time", "none", "both") 
re.use = "both"

# whether the time-level random effects is ar1 process or jsut multilevel (independent)
ar1.use = TRUE

#factor numbers
r.use = 10

niter.use = 15000 # number of mcmc draws
burn.use = 5000   # burn-in draws 
xlasso.use = 1    ## whether to shrink constant coefs (1 = TRUE, 0 = FALSE)
zlasso.use = 1    ## whether to shrink unit-level random coefs (1 = TRUE, 0 = FALSE)
alasso.use = 1    ## whether to shrink time-level coefs (1 = TRUE, 0 = FALSE)
flasso.use = 1    ## whether to shrink factor loadings (1 = TRUE, 0 = FALSE)

## parameters for hyper prior shrink on beta (diffuse hyper priors)
a1.use = 0.001
a2.use = 0.001 

## parameters for hyper prior shrink on alpha_i
b1.use = 0.001
b2.use = 0.001 

## parameters for hyper prior shrink on xi_t
c1.use = 0.001
c2.use = 0.001 

## parameters for hyper prior shrink on factor terms
p1.use = 0.001
p2.use = 0.001

##seed
seed.use = 1234

#Index
index.use = c("iso2c", "year")

#---------------------------------------------------
#Function to bpcausal by group
#---------------------------------------------------
bpcausal.group = function(dv.list,
                          D.use,
                          X.use, Z.use, A.use,
                          df.use,
                          index = index.use,
                          re = re.use,
                          ar1 = ar1.use,
                          r = r.use,
                          niter = niter.use,
                          burn = burn.use,
                          xlasso = xlasso.use,
                          zlasso = zlasso.use,
                          alasso = alasso.use,
                          flasso = flasso.use,
                          a1 = a1.use,
                          a2 = a2.use,
                          b1 = b1.use,
                          b2 = b2.use,
                          c1 = c1.use,
                          c2 = c2.use,
                          p1 = p1.use,
                          p2 = p2.use,
                          seed = seed.use
){
  out.list = list()
  for (i in 1:length(dv.list)){
    set.seed(seed)
    Y.use = dv.list[i]
    out.list[[i]] = bpCausal(data = df.use,
                             Yname = Y.use,
                             Dname = D.use,
                             Xname = X.use,
                             Zname = Z.use,
                             Aname = A.use,
                             index = index.use,
                             re = re,
                             ar1 = ar1,
                             r = r,
                             niter = niter,
                             burn = burn,
                             xlasso = xlasso,
                             zlasso = zlasso,
                             alasso = alasso,
                             flasso = flasso,
                             a1 = a1,
                             a2 = a2,
                             b1 = b1,
                             b2 = b2,
                             c1 = c1,
                             c2 = c2,
                             p1 = p1,
                             p2 = p2)
    print(i)
  }
  names(out.list) = dv.list
  return(out.list)
}

######################################
#Estimate using feglm with clustered SE
######################################
feglm.use = function(eq, df){
  fit = feglm(eq, data = df, family = binomial())
  se = summary(fit, type = "sandwich")
  return(list(fit = fit,
              se = se))
}